"""
市场状态检测模块
用于检测和分类不同的市场状态（趋势、震荡等）
"""

import numpy as np
import pandas as pd
from typing import Dict, List, Tuple, Optional
from datetime import datetime


class MarketRegimeDetector:
    """市场状态检测器"""
    
    def __init__(self, window_size: int = 20):
        """
        初始化市场状态检测器
        
        Args:
            window_size: 分析窗口大小
        """
        self.window_size = window_size
    
    def detect_market_regime(self, prices: pd.Series) -> Dict[str, any]:
        """
        检测市场状态
        
        Args:
            prices: 价格序列
            
        Returns:
            Dict: 市场状态信息
        """
        if len(prices) < self.window_size:
            return {
                'regime': 'unknown',
                'confidence': 0.0,
                'volatility': 0.0,
                'trend_strength': 0.0
            }
        
        # 计算波动率
        volatility = self._calculate_volatility(prices)
        
        # 计算趋势强度
        trend_strength = self._calculate_trend_strength(prices)
        
        # 检测市场状态
        regime, confidence = self._classify_regime(volatility, trend_strength)
        
        return {
            'regime': regime,
            'confidence': confidence,
            'volatility': volatility,
            'trend_strength': trend_strength
        }
    
    def _calculate_volatility(self, prices: pd.Series) -> float:
        """
        计算波动率
        
        Args:
            prices: 价格序列
            
        Returns:
            float: 波动率
        """
        returns = prices.pct_change().dropna()
        return np.std(returns) * np.sqrt(252)  # 年化波动率
    
    def _calculate_trend_strength(self, prices: pd.Series) -> float:
        """
        计算趋势强度
        
        Args:
            prices: 价格序列
            
        Returns:
            float: 趋势强度 (-1 到 1)
        """
        if len(prices) < 2:
            return 0.0
        
        # 简单线性回归计算趋势
        x = np.arange(len(prices))
        y = prices.values
        
        # 计算斜率
        slope, _ = np.polyfit(x, y, 1)
        
        # 标准化趋势强度
        price_range = prices.max() - prices.min()
        if price_range == 0:
            return 0.0
            
        normalized_slope = slope / (price_range / len(prices))
        trend_strength = np.clip(normalized_slope, -1, 1)
        
        return trend_strength
    
    def _classify_regime(self, volatility: float, trend_strength: float) -> Tuple[str, float]:
        """
        分类市场状态
        
        Args:
            volatility: 波动率
            trend_strength: 趋势强度
            
        Returns:
            Tuple: (市场状态, 置信度)
        """
        # 根据波动率和趋势强度分类
        if abs(trend_strength) > 0.5 and volatility < 0.2:
            regime = 'strong_trend'  # 强趋势
            confidence = min(1.0, abs(trend_strength))
        elif abs(trend_strength) > 0.2 and volatility < 0.3:
            regime = 'trend'  # 趋势
            confidence = min(0.8, abs(trend_strength))
        elif abs(trend_strength) <= 0.2 and volatility > 0.25:
            regime = 'volatile'  # 震荡
            confidence = min(1.0, volatility)
        else:
            regime = 'mixed'  # 混合
            confidence = 0.5
            
        return regime, confidence


def main():
    """主函数 - 市场状态检测演示"""
    print("市场状态检测模块演示")
    print("="*30)
    
    # 创建示例数据
    np.random.seed(42)
    dates = pd.date_range('2023-01-01', periods=100, freq='D')
    prices = pd.Series(
        [100 * (1 + np.random.normal(0, 0.02))**i for i in range(100)], 
        index=dates
    )
    
    # 创建检测器
    detector = MarketRegimeDetector(window_size=20)
    
    # 检测市场状态
    regime_info = detector.detect_market_regime(prices)
    
    print(f"\n市场状态检测结果:")
    print(f"状态: {regime_info['regime']}")
    print(f"置信度: {regime_info['confidence']:.2f}")
    print(f"波动率: {regime_info['volatility']:.2%}")
    print(f"趋势强度: {regime_info['trend_strength']:.2f}")

if __name__ == "__main__":
    main()